import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt
import sqlite3
import numpy as np



conn = sqlite3.connect('./sqlData.db')

def get_data(bs_code='sz.000001',start_date='2011-01-01',end_date='2020-05-20',conn=''):
    rs = bs.query_history_k_data_plus(code=bs_code,fields='date,code,open,high,low,close,preclose,pctChg,volume,amount,turn,isST',
                                 start_date=start_date,end_date=end_date,frequency='d',adjustflag='2')
    print('error_code:'+rs.error_code)
    print('error_msg:'+rs.error_msg)
    data_list = []
    while rs.next() & (rs.error_code=='0'):
        data_list.append(rs.get_row_data())
    rst = pd.DataFrame(data_list,columns=rs.fields)
    # print(type(rst['pctChg']))
    isDayLimited = len([i for i in rst['pctChg']>'9.93' if i == True])>0
    if isDayLimited:

        rst.to_sql('t'+bs_code[-6:],conn)


def computeKDJ(bs_code='sh.600001',start_date='',end_date=''):
    print(bs_code,start_date,end_date)
    login_rst = bs.login(user_id='anonymous',password='123456')
    rs = bs.query_history_k_data_plus(code=bs_code,fields='date,code,high,close,low,tradeStatus',
                                      start_date=start_date,end_date=end_date,frequency='d',adjustflag='2')
    rst_list = []
    while (rs.error_code=='0') & rs.next():
        rst_list.append(rs.get_row_data())
    data_init = pd.DataFrame(rst_list,columns=rs.fields)
    # print(data_init)
    data_status = data_init[data_init['tradeStatus'] == '1']
    low = data_status['low'].astype(float)
    # print(data_status)
    del data_status['low']
    data_status.insert(0,'low',low)
    high = data_status['high'].astype(float)
    del data_status['high']
    data_status.insert(0,'high',high)
    close = data_status['close'].astype(float)
    del data_status['close']
    data_status.insert(0,'close',close)
    low_list = data_status['low'].rolling(window=9).min()
    print(low_list)
    high_list = data_status['high'].rolling(window=9).max()
    rsv = (data_status['close'] - low_list)/(high_list - low_list)*100
    df_data = pd.DataFrame()
    df_data['High'] = data_status['high']
    df_data['Low'] = data_status['low']
    df_data['Close'] = data_status['close']
    df_data['K'] = rsv.ewm(com=2).mean()
    df_data['D'] = df_data['K'].ewm(com=2).mean()
    df_data['J'] = 3*df_data['K'] - 2*df_data['D']
    df_data.index = data_status['date'].values
    df_data.index.name = 'date'
    df_data = df_data.dropna()
    # print(df_data)
    df_data['KDJ_金叉死叉'] = ''
    # kdj_position = df_data['K'] > df_data['D']
    kdj_position = df_data['J']< 2.1
    kdj_pos = df_data['J']> 94.5
    df_data.loc[kdj_position[(kdj_position == True) & (kdj_position.shift() == False)].index, 'KDJ_金叉死叉'] = '金叉'
    df_data.loc[kdj_pos[(kdj_pos == True) & (kdj_pos.shift() == False)].index, 'KDJ_金叉死叉'] = '死叉'
    df_data.plot(title='KDJ')

    plt.show()
    bs.logout()
    return (df_data)


def get_industry_data():
    # get all industry data save in table industry_data
    rs = bs.query_stock_industry()
    industry_list = []
    while rs.next():
        industry_list.append(rs.get_row_data())
    rst = pd.DataFrame(industry_list,columns=rs.fields)
    rst.to_sql('industry_data',conn)


def get_profit_data(code,year,quarter):
    # return every quarter profit data
    rs = bs.query_profit_data(code,year,quarter)
    profit_list = []
    while rs.next():
        profit_list.append(rs.get_row_data())
    rst = pd.DataFrame(profit_list,columns=rs.fields)
    return rst


def update_month_data(start,end,conn):
    # 获取满足条件的股票近一个月的数据
    cur = conn.cursor()
    cur.execute("select name from sqlite_master where name like 't%'")
    tables = cur.fetchall()
    if len(tables)>0:
        for row in tables:
            cur.execute("drop table %s"%row[0])
        conn.commit()
    else:
        cur.execute("select code from industry_data where industry in('计算机','电子','通信','传媒')")
        for row in cur:
            code = row[0]
            get_data(code,start,end,conn)


def get_growth_data(code,year,quarter):
    growth_list = []
    df_fields = ["证券代码","发布日期","统计日期","净资产增长率","总资产增长率","净利润增长率","每股收益增长率","母公司增长率",
                   "净益率","销售净利率","毛利率","净利润","每股收益","主营收入","总股本","流通股"]
    for yer in range(year,2020):
        rs_growth=bs.query_growth_data(code=code,year=yer,quarter=quarter)
        rs_profit = bs.query_profit_data(code=code,year=yer,quarter=quarter)
        while (rs_growth.error_code == "0") & rs_growth.next():
            growth_list.append(rs_growth.get_row_data()+rs_profit.get_row_data()[3:])

    rst_growth = pd.DataFrame(growth_list,columns=df_fields)
    print(growth_list)
    rst_growth = pd.DataFrame(growth_list,columns=df_fields)
    rst_growth['净利润增长率'] =rst_growth['净利润增长率'].astype(float)
    rst_growth['净利润']=rst_growth['净利润'].astype(float)/100000000
    # rst_growth['主营收入']=rst_growth['主营收入'].astype(float)/100000000

    rst_growth.to_csv('./'+code+'.csv',encoding='gbk')
    return rst_growth


def get_good_profit_stock(startyear, quarter):
    query_str ="select code, code_name from industry_data where industry in('电子') order by code DESC"
    cur = conn.cursor()
    cur.execute(query_str)
    for row in cur:
        code,code_name = row[0],row[1]
        profit_df = get_growth_data(code, startyear, quarter)
        if list(profit_df['净利润增长率'])[-2:] > [0.1] and list(profit_df['净利润'])[-2:] > 0:
            profit_df.to_csv('./'+code_name[-4:]+'.csv',encoding='gbk')


def main():
    lg = bs.login()
    # get_data(conn=conn)
    # df = computeKDJ('sh.603636','2019-05-12','2020-05-28')
    # df.to_excel("D:/kdj603636.xlsx")
    # get_industry_data()
    # update_month_data('2020-05-01','2020-06-09',conn)
    # get_data('sh.600446','','',)
    get_growth_data("sz.002218",2010,4)
    # get_good_profit_stock(2010,4)
    bs.logout()


if __name__ == '__main__':
    main()
